use data_structure::prelude::*;
use std::ops::{Add, Sub};

struct Polynomial(Ring<DNode<(i64, i64)>>);

impl Polynomial {
    #[inline]
    pub fn create(c: Vec<i64>, e: Vec<i64>) -> Self {
        assert!(c.len() == e.len());
        let mut ring: Ring<DNode<(i64, i64)>> = Ring::new();
        c.into_iter()
            .zip(e.into_iter())
            .for_each(|params| ring.push(params));
        ring.roll(1); // roll to head
        Polynomial(ring)
    }

    #[inline]
    pub fn differentiate(&mut self) {
        for idx in 0..self.0.size() {
            if let Some((c, e)) = self.0.at(idx) {
                match e == 0 {
                    false => self.0.update(idx, (c * e, e - 1)),
                    true => {
                        self.0.delete(idx);
                    }
                }
            }
        }
    }

    #[inline]
    fn solve(&self, x: i64) -> i64 {
        let mut accum = 0i64;
        for idx in (0..self.0.size()).rev() {
            if let Some((c, e)) = self.0.at(idx) {
                assert!(e >= 0);
                accum += c * x.pow(e as u32);
            }
        }
        accum
    }
}

impl Add for Polynomial {
    type Output = Polynomial;

    #[inline]
    fn add(self, rhs: Polynomial) -> Self::Output {
        let mut ring: Ring<DNode<(i64, i64)>> = Ring::new();
        let (size_1, size_2) = (self.0.size(), rhs.0.size());
        let (mut idx_1, mut params_1) = (0, self.0.at(0));
        let (mut idx_2, mut params_2) = (0, rhs.0.at(0));
        while idx_1 < size_1 || idx_2 < size_2 {
            let (c_1, e_1) = match params_1 {
                Some((c, e)) => (c, e),
                None => (0, 0),
            };
            let (c_2, e_2) = match params_2 {
                Some((c, e)) => (c, e),
                None => (0, 0),
            };
            if e_1 < e_2 {
                ring.push((c_2, e_2));
                idx_2 += 1;
                params_2 = rhs.0.at(idx_2);
            } else if e_1 > e_2 {
                ring.push((c_1, e_1));
                idx_1 += 1;
                params_1 = self.0.at(idx_1);
            } else {
                if c_1 + c_2 != 0 {
                    ring.push((c_1 + c_2, e_2));
                }
                idx_1 += 1;
                idx_2 += 1;
                params_1 = self.0.at(idx_1);
                params_2 = rhs.0.at(idx_2);
            }
        }
        Polynomial(ring)
    }
}

impl Sub for Polynomial {
    type Output = Polynomial;

    #[inline]
    fn sub(self, rhs: Polynomial) -> Self::Output {
        let mut ring: Ring<DNode<(i64, i64)>> = Ring::new();
        let (size_1, size_2) = (self.0.size(), rhs.0.size());
        let (mut idx_1, mut params_1) = (0, self.0.at(0));
        let (mut idx_2, mut params_2) = (0, rhs.0.at(0));
        while idx_1 < size_1 || idx_2 < size_2 {
            let (c_1, e_1) = match params_1 {
                Some((c, e)) => (c, e),
                None => (0, 0),
            };
            let (c_2, e_2) = match params_2 {
                Some((c, e)) => (c, e),
                None => (0, 0),
            };
            if e_1 < e_2 {
                ring.push((-c_2, e_2));
                idx_2 += 1;
                params_2 = rhs.0.at(idx_2);
            } else if e_1 > e_2 {
                ring.push((c_1, e_1));
                idx_1 += 1;
                params_1 = self.0.at(idx_1);
            } else {
                if c_1 != c_2 {
                    ring.push((c_1 - c_2, e_2));
                }
                idx_1 += 1;
                idx_2 += 1;
                params_1 = self.0.at(idx_1);
                params_2 = rhs.0.at(idx_2);
            }
        }
        Polynomial(ring)
    }
}

#[test]
fn poly_calculate() {
    let mut poly = Polynomial::create(vec![3, 2, 1], vec![2, 1, 0]);
    assert_eq!(poly.solve(2), 17);
    let sum_of_power = Polynomial::create(vec![1, 1, 1, 1], vec![3, 2, 1, 0]);
    poly = poly + sum_of_power;
    poly.differentiate();
    assert_eq!(poly.solve(1), 14);
}
